package com.itcast.peft.config;

import java.util.Optional;

import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.model.SimpleApiKey;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.util.StringUtils;

import lombok.AllArgsConstructor;
import lombok.Data;

@Configuration
public class MultipleModelConfig {

    @Bean
    public ChatModelHolder healthGenModelHolder() {
        var cfg = healthGenModelConfig();
        return new ChatModelHolder(buildByConfig(cfg), cfg);
    }
    
    @Bean
    @ConfigurationProperties(prefix = "health-gen")
    public BaseModelConfig healthGenModelConfig() {
        return new BaseModelConfig();
    }
    
    @Bean
    public ChatModelHolder healthOptModelHolder() {
        var cfg = healthOptModelConfig();
        return new ChatModelHolder(buildByConfig(cfg), cfg);
    }
    
    @Bean
    @ConfigurationProperties(prefix = "health-opt")
    public BaseModelConfig healthOptModelConfig() {
        return new BaseModelConfig();
    }
    
    @Bean
    public ChatModelHolder userprefEvalModelHolder() {
        var cfg = userprefEvalModelConfig();
        return new ChatModelHolder(buildByConfig(cfg), cfg);
    }
    
    @Bean
    @ConfigurationProperties(prefix = "userpref-eval")
    public BaseModelConfig userprefEvalModelConfig() {
        return new BaseModelConfig();
    }
    
    @Bean
    public ChatModelHolder userprefOptModelHolder() {
        var cfg = userprefOptModelConfig();
        return new ChatModelHolder(buildByConfig(cfg), cfg);
    }
    
    @Bean
    @ConfigurationProperties(prefix = "userpref-opt")
    public BaseModelConfig userprefOptModelConfig() {
        return new BaseModelConfig();
    }
    
    @Bean
    public ChatModelHolder userprefGenerationModelHolder() {
        var cfg = userprefGenerationModelConfig();
        return new ChatModelHolder(buildByConfig(cfg), cfg);
    }
    
    @Bean
    @ConfigurationProperties(prefix = "userpref-generation")
    public BaseModelConfig userprefGenerationModelConfig() {
        return new BaseModelConfig();
    }
    
    @Bean
    public ChatModelHolder distillStudentModelHolder() {
        var cfg = studentModelConfig();
        return new ChatModelHolder(buildByConfig(cfg), cfg);
    }

    @Bean
    @ConfigurationProperties(prefix = "distill-student")
    public BaseModelConfig studentModelConfig() {
        return new BaseModelConfig();
    }

    @Bean
    public ChatModelHolder distillEvalModelHolder() {
        var cfg = evalModelConfig();
        return new ChatModelHolder(buildByConfig(cfg), cfg);
    }

    @Bean
    @ConfigurationProperties(prefix = "distill-eval")
    public BaseModelConfig evalModelConfig() {
        return new BaseModelConfig();
    }

    @Bean
    public ChatModelHolder distillQuestionModelHolder() {
        var cfg = questionModelConfig();
        return new ChatModelHolder(buildByConfig(cfg), cfg);
    }

    @Bean
    @ConfigurationProperties(prefix = "distill-question")
    public BaseModelConfig questionModelConfig() {
        return new BaseModelConfig();
    }

    @Bean
    public ChatModelHolder distillTeacherModelHolder() {
        var cfg = teacherModelConfig();
        return new ChatModelHolder(buildByConfig(cfg), cfg);
    }

    @Bean
    @ConfigurationProperties(prefix = "distill-teacher")
    public BaseModelConfig teacherModelConfig() {
        return new BaseModelConfig();
    }

    private ChatModel buildByConfig(BaseModelConfig cfg) {
        var bd = OpenAiApi.builder().baseUrl(cfg.getBaseUrl()).apiKey(new SimpleApiKey(cfg.getApiKey()));

        Optional.ofNullable(cfg.getCompletionsPath()).filter(StringUtils::hasText).ifPresent(bd::completionsPath);

        var option = OpenAiChatOptions.builder().model(cfg.getModelName()).build();
        return OpenAiChatModel.builder().openAiApi(bd.build()).defaultOptions(option).build();
    }

    @Data
    public static class BaseModelConfig {
        String baseUrl;
        String apiKey;
        String modelName;
        String completionsPath;
    }

    @Data
    @AllArgsConstructor
    public static class ChatModelHolder {
        private final ChatModel model;
        private final BaseModelConfig config;
    }
}
